#include <string.h>
#include <openssl/sm3.h>
#include "crypto.h"
#include "session_data.h"
#include "ovmf_hash_table.h"
#include "vmsa.h"
#include "utils.h"

#define ALIGN_UP(size, align) (((size) + (align)-1) & (~((align)-1)))

/*
 *  4 type of data structures the VM digest will cover
 *  full measurement means bios+kernel+cmdline+initrd measurement
 *
 *  type 1
 *  CSV guest bios measurement support
 *  |    bios file                              |
 *
 *  type 2
 *  CSV guest bios/kernel/initrd/cmdline measurement support
 *  |    bios file                              |
 *  |    hash table                             |
 *
 *  type 3
 *  CSV3 guest bios/vcpu measurement support
 *  |    csv3_measure_structure_of(bios)        |
 *  |    bios file                              |
 *  |    csv3_measure_structure_of(vcpu0)       |
 *  |    vcpu0 VMSA                             |
 *  |    csv3_measure_structure_of(vcpu1)       |
 *  |    vcpu1 VMSA                             |
 *  |    csv3_measure_structure_of(vcpuN)       |
 *  |    vcpuN VMSA                             |
 *
 *  type 4
 *  CSV3 guest bios/kernel/initrd/cmdline/vcpu measurement support
 *  |    csv3_measure_structure_of(bios)        |
 *  |    bios file                              |
 *  |    csv3_measure_structure_of(hash table)  |
 *  |    hash table                             |
 *  |    csv3_measure_structure_of(vcpu0)       |
 *  |    vcpu0 VMSA                             |
 *  |    csv3_measure_structure_of(vcpu1)       |
 *  |    vcpu1 VMSA                             |
 *  |    csv3_measure_structure_of(vcpuN)       |
 *  |    vcpuN VMSA                             |
 *
 */

/* type 1 */
/**
 * @brief  calculate the CSV VM bios digest
 *
 * @param  [in]   bios_file  The bios file name
 * @param  [out]  digest     The output digest
 *
 *
 * @returns 0 for success, -1 for failure
 */
int calculate_csv_bios_digest(char *bios_file, uint8_t digest[HASH_BLOCK_LEN])
{
    int ret = -1;

    uint32_t    bios_file_size          = 0;
    uint32_t    buffer_size             = 0;
    uint8_t*    launch_update_buffer    = NULL;

    bios_file_size = get_file_length(bios_file);

    buffer_size = bios_file_size;

    launch_update_buffer = (uint8_t*)malloc(buffer_size);
    if (!launch_update_buffer) {
        printf("fail to alloc %u bytes\n", buffer_size);
        goto finish;
    }

    memset(launch_update_buffer, 0, buffer_size);

    // fill in bios
    read_file(bios_file, launch_update_buffer, bios_file_size);

    sm3(launch_update_buffer, buffer_size, (uint8_t*)digest);

    ret = 0;

finish:
    if (launch_update_buffer) {
        free(launch_update_buffer);
    }

    return ret;
}

/* type 2 */
/**
 * @brief  calculate the full CSV VM launch digest
 *
 * @param  [in]   bios_file    The bios file name
 * @param  [in]   kernel_file  The kernel file name
 * @param  [in]   cmdline      The command line string
 * @param  [in]   initrd_file  The initrd file name
 * @param  [out]  digest       The output digest
 *
 *
 * @returns 0 for success, -1 for failure
 */
int calculate_csv_full_digest(char *bios_file, char *kernel_file,
                              char *cmdline, char *initrd_file,
                              uint8_t digest[HASH_BLOCK_LEN])
{
    int ret = -1;

    uint32_t     bios_file_size          = 0;
    uint32_t     buffer_size             = 0;
    uint8_t*     launch_update_buffer    = NULL;
    uint32_t     hash_table_size_aligned = 0;
    uint8_t*     ptr;
    hash_table   *ht;

    hash_table_size_aligned = ALIGN_UP(sizeof(hash_table), 16);

    bios_file_size = get_file_length(bios_file);

    buffer_size += bios_file_size;
    buffer_size += hash_table_size_aligned;

    launch_update_buffer = (uint8_t*)malloc(buffer_size);
    if (!launch_update_buffer) {
        printf("fail to alloc %u bytes\n", buffer_size);
        goto finish;
    }
    memset(launch_update_buffer, 0, buffer_size);

    ptr = launch_update_buffer;

    // fill in bios
    read_file(bios_file, ptr, bios_file_size);

    ptr += bios_file_size;

    ht = (hash_table *)ptr;

    if (fill_csv_hash_table(ht, kernel_file, cmdline, initrd_file)) {
        printf("fill sev hash table failed!\n");
        goto finish;
    }
    if (hash_table_size_aligned != ht->len)
        memset(ht->padding, 0, hash_table_size_aligned - ht->len);

    sm3(launch_update_buffer, buffer_size, (uint8_t*)digest);

    ret = 0;

finish:
    if (launch_update_buffer) {
        free(launch_update_buffer);
    }

    return ret;

}

void fill_csv3_bios_structure(uint8_t *ptr, char *bios_file)
{
    uint32_t       bios_file_size  = 0;
    csv3_measure_update_extra_data_t* extra_data_ptr = NULL;

    bios_file_size = get_file_length(bios_file);

    extra_data_ptr = (csv3_measure_update_extra_data_t *)ptr;
    extra_data_ptr->command_id = (uint32_t)CSV3_LAUNCH_ENCRYPT_DATA;
    extra_data_ptr->gpa = 0x100000000ull - bios_file_size;
    extra_data_ptr->data_size = bios_file_size;
    ptr += sizeof(csv3_measure_update_extra_data_t);

    // fill in bios
    read_file(bios_file, ptr, bios_file_size);
}

int  fill_csv3_hash_table_structure(uint8_t *ptr,
                                    uint8_t *bios_ptr, uint32_t bios_file_size,
                                    char *kernel_file, char *cmdline,
                                    char *initrd_file)
{
    int          ret          = -1;
    uint32_t     gpa;
    csv3_measure_update_extra_data_t* extra_data_ptr = NULL;
    uint32_t     hash_table_size_aligned = 0;
    hash_table   *ht;

    hash_table_size_aligned = ALIGN_UP(sizeof(hash_table), 16);

    if (ovmf_find_hash_table_desc(bios_ptr, bios_file_size, &gpa)) {
       printf("hashes table area failed!\n");
       goto finish;
    }

    extra_data_ptr = (csv3_measure_update_extra_data_t *)ptr;
    extra_data_ptr->command_id = (uint32_t)CSV3_LAUNCH_ENCRYPT_DATA;
    extra_data_ptr->gpa = (uint64_t)gpa;
    extra_data_ptr->data_size = hash_table_size_aligned;
    ptr += sizeof(csv3_measure_update_extra_data_t);

    ht = (hash_table *)ptr;

    if (fill_csv_hash_table(ht, kernel_file, cmdline, initrd_file)) {
        printf("fill sev hash table failed!\n");
        goto finish;
    }

    if (hash_table_size_aligned != ht->len) {
        memset(ht->padding, 0, hash_table_size_aligned - ht->len);
    }

    ret = 0;
finish:

    return ret;
}

int fill_csv3_vcpu_structure(uint8_t *ptr,
                             uint8_t *bios_ptr, uint32_t bios_file_size,
                             char* vcpu_model, unsigned long vcpu_num)
{
    int ret = -1;
    int vcpu_id;
    csv3_measure_update_extra_data_t* extra_data_ptr = NULL;
    uint32_t     gpa;


    if (ovmf_find_reset_vector(bios_ptr, bios_file_size, &gpa)) {
        printf("find reset vector area failed!\n");
        goto finish;
     }

    for (vcpu_id = 0; vcpu_id < vcpu_num; vcpu_id++) {
        extra_data_ptr = (csv3_measure_update_extra_data_t *)ptr;
        extra_data_ptr->command_id = (uint32_t)CSV3_LAUNCH_ENCRYPT_VMCB;
        extra_data_ptr->vcpu_id = (uint64_t)vcpu_id;
        extra_data_ptr->data_size = VMSA_SIZE;
        ptr += sizeof(csv3_measure_update_extra_data_t);
        vcpu_reset((void *)ptr, (size_t)VMSA_SIZE,
                                  vcpu_model,
                                  vcpu_id ? gpa : 0, vcpu_id ? true : false);
        ptr += VMSA_SIZE;
    }

    ret = 0;
finish:

    return ret;
}

/* type 3 */
/**
 * @brief  calculate the CSV3 VM bios+vcpu digest
 *
 * @param  [in]   bios_file   The bios file name
 * @param  [in]   vcpu_model  The vcpu model
 * @param  [in]   vcpu_num    The vcpu number
 * @param  [out]  digest      The output digest
 *
 *
 * @returns 0 for success, -1 for failure
 */
int calculate_csv3_bios_digest(char *bios_file,
                              char* vcpu_model, unsigned long vcpu_num,
                              uint8_t digest[HASH_BLOCK_LEN])
{
    int ret = -1;

    uint32_t                          bios_file_size          = 0;
    uint32_t                          buffer_size             = 0;
    uint8_t*                          launch_update_buffer    = NULL;
    uint8_t*                          ptr;
    uint8_t *bios_ptr;

    bios_file_size = get_file_length(bios_file);

    buffer_size += sizeof(csv3_measure_update_extra_data_t);
    buffer_size += bios_file_size;
    buffer_size += vcpu_num * (sizeof(csv3_measure_update_extra_data_t) + VMSA_SIZE);

    launch_update_buffer = (uint8_t*)malloc(buffer_size);
    if (!launch_update_buffer) {
        printf("fail to alloc %u bytes\n", buffer_size);
        goto finish;
    }

    memset(launch_update_buffer, 0, buffer_size);

    fill_csv3_bios_structure(launch_update_buffer, bios_file);
    bios_ptr = launch_update_buffer + sizeof(csv3_measure_update_extra_data_t);

    /* move pointer to vcpu area */
    ptr = launch_update_buffer + sizeof(csv3_measure_update_extra_data_t)
          + bios_file_size;

    if (fill_csv3_vcpu_structure(ptr, bios_ptr, bios_file_size,
                              vcpu_model, vcpu_num) != 0) {
        printf("fail to fill csv3 vcpu structure\n");
        goto finish;
    }

    sm3(launch_update_buffer, buffer_size, (uint8_t*)digest);

    ret = 0;

finish:
    if (launch_update_buffer) {
        free(launch_update_buffer);
    }

    return ret;
}

/* type 4 */
/**
 * @brief  calculate the full CSV3 VM launch digest
 *
 * @param  [in]   bios_file    The bios file name
 * @param  [in]   kernel_file  The kernel file name
 * @param  [in]   cmdline      The command line string
 * @param  [in]   initrd_file  The initrd file name
 * @param  [in]   vcpu_model  The vcpu model
 * @param  [in]   vcpu_num    The vcpu number
 * @param  [out]  digest       The output digest
 *
 *
 * @returns 0 for success, -1 for failure
 */
int calculate_csv3_full_digest(char *bios_file, char *kernel_file,
                              char *cmdline, char *initrd_file,
                              char* vcpu_model, unsigned long vcpu_num,
                              uint8_t digest[HASH_BLOCK_LEN])
{
    int ret = -1;

    uint32_t      bios_file_size          = 0;
    uint32_t      buffer_size             = 0;
    uint8_t*      launch_update_buffer    = NULL;
    uint8_t*      ptr;
    uint8_t *bios_ptr;
    uint32_t     hash_table_size_aligned = 0;

    hash_table_size_aligned = ALIGN_UP(sizeof(hash_table), 16);

    bios_file_size = get_file_length(bios_file);

    buffer_size += sizeof(csv3_measure_update_extra_data_t);
    buffer_size += bios_file_size;
    buffer_size += sizeof(csv3_measure_update_extra_data_t);
    buffer_size += hash_table_size_aligned;
    buffer_size += vcpu_num * (sizeof(csv3_measure_update_extra_data_t) + VMSA_SIZE);

    launch_update_buffer = (uint8_t*)malloc(buffer_size);
    if (!launch_update_buffer) {
        printf("fail to alloc %u bytes\n", buffer_size);
        goto finish;
    }

    memset(launch_update_buffer, 0, buffer_size);

    fill_csv3_bios_structure(launch_update_buffer, bios_file);
    bios_ptr = launch_update_buffer + sizeof(csv3_measure_update_extra_data_t);

    /* pointer to hash area */
    ptr = launch_update_buffer + sizeof(csv3_measure_update_extra_data_t)
          + bios_file_size;
    if (fill_csv3_hash_table_structure(ptr, bios_ptr, bios_file_size,
                                   kernel_file, cmdline, initrd_file) != 0) {
        printf("fail to fill csv3 hash table structure\n");
        goto finish;
    }

    /* pointer to vcpu area */
    ptr +=  sizeof(csv3_measure_update_extra_data_t) + hash_table_size_aligned;

    if (fill_csv3_vcpu_structure(ptr, bios_ptr, bios_file_size,
                             vcpu_model, vcpu_num) != 0) {
        printf("fail to fill csv3 vcpu structure\n");
        goto finish;
    }

    sm3(launch_update_buffer, buffer_size, (uint8_t*)digest);

    ret = 0;

finish:
    if (launch_update_buffer) {
        free(launch_update_buffer);
    }

    return ret;

}

/**
 * @brief  calculate the digest of the VM
 *
 * @param  [in]  policy       The guest policy
 * @param  [in]  bios_file    The VM bios file name
 * @param  [in]  kernel_file  The VM kernel file name
 * @param  [in]  cmdline      The VM kernel boot command line
 * @param  [in]  initrd_file  The VM initrd file name
 * @param  [in]  vcpu_model   The VM vcpu model
 * @param  [in]  vcpu_num     The VM vcpu number
 * @param  [out] digest       The VM digest
 *
 * @returns 0 for success, -1 for failure
 */
int  calculate_vm_digest(guest_policy policy,
                         char *bios_file, char *kernel_file,
                         char *cmdline, char *initrd_file,
                         char* vcpu_model, unsigned long vcpu_num,
                         uint8_t digest[HASH_BLOCK_LEN])
{
    int ret = -1;

    uint32_t                          file_size          = 0;
    uint32_t                          hash_table_size_aligned = 0;

    if (policy.csv3 == 1 && (vcpu_num == 0 || vcpu_model == NULL)) {
        printf("must specify -cpu and -smp for CSV3 guest\n");
        return ret;
    }

    file_size = get_file_length(bios_file);
    if (file_size == 0) {
        printf("invalid bios length 0\n");
        goto finish;
    }

    /*
     * If there is kernel file,
     * need hash_table for full process measuremenet.
      * */
    if (kernel_file && strcmp(kernel_file, "")!=0) {
        if (initrd_file == NULL) {
            printf("kernel file is provided, but initrd is NULL\n");
            goto finish;
        }

        file_size = get_file_length(kernel_file);
        if (file_size == 0) {
            printf("invalid kernel file length\n");
            goto finish;
        }

        file_size = get_file_length(initrd_file);
        if (file_size == 0) {
            printf("invalid initrd file length\n");
            goto finish;
        }

        hash_table_size_aligned = ALIGN_UP(sizeof(hash_table), 16);
    }

    if (policy.csv3 == 0) {
        if (hash_table_size_aligned == 0) {
            ret = calculate_csv_bios_digest(bios_file, digest);
        } else {
            ret = calculate_csv_full_digest(bios_file, kernel_file,
                                      cmdline, initrd_file, digest);
        }
    } else {
        if (hash_table_size_aligned == 0) {
            ret = calculate_csv3_bios_digest(bios_file, vcpu_model,vcpu_num,
                                       digest);
        } else {
            ret = calculate_csv3_full_digest(bios_file, kernel_file,
                                       cmdline, initrd_file,
                                       vcpu_model,vcpu_num,
                                       digest);
        }
    }

finish:
    return ret;
}
